/* AT25DF SPI Flash driver */

/* Copyright (C) 2011 David Zanetti
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; version 2 of the License.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License along
 *  with this program; if not, write to the Free Software Foundation, Inc.,
 *  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */

#include <avr/io.h>
#include <avr/pgmspace.h>
#include <stdio.h>
#include <string.h>
#include <avr/eeprom.h>
#include "global.h"
#include <util/delay.h>
#include "at25df.h"
#include "spi.h"
#include "serial.h"

#define AT25_READ_ARRAY 0x0b
#define AT25_READ_ARRAY_SLOW 0x03
#define AT25_BLOCK_ERASE_4K 0x20
#define AT25_BLOCK_ERASE_32K 0x52
#define AT25_BLOCK_ERASE_64K 0xd8
#define AT25_CHIP_ERASE_A 0x60
#define AT25_CHIP_ERASE_B 0xc7
#define AT25_BYTEPAGE_PROGRAM 0x02
#define AT25_SEQUENCE_PROGRAM_A 0xad
#define AT25_SEQUENCE_PROGRAM_B 0xaf
#define AT25_WRITE_ENABLE 0x06
#define AT25_WRITE_DISABLE 0x04
#define AT25_PROTECT_SECTOR 0x36
#define AT25_UNPROTECT_SECTOR 0x39
#define AT25_READ_SECTOR_PROTECT_REG 0x3c
#define AT25_READ_STATUS 0x05
#define AT25_WRITE_STATUS 0x01
#define AT25_READ_MANUDEVID 0x9f
#define AT25_DEEP_POWERDOWN 0xb9
#define AT25_RESUME_FROM_DEEP 0xab

#define AT25_STATUS_SPRL 0x80
#define AT25_STATUS_SPM 0x40
#define AT25_STATUS_EPE 0x20
#define AT25_STATUS_WPP 0x10
#define AT25_STATUS_SWP1 0x08
#define AT25_STATUS_SWP0 0x04
#define AT25_STATUS_WEL 0x02
#define AT25_STATUS_BSY 0x01

#define FLASH_CS_PORT PORTE
#define FLASH_CS_PIN PIN0_bm

#define SPI_START FLASH_CS_PORT.OUTCLR = FLASH_CS_PIN;
#define SPI_END FLASH_CS_PORT.OUTSET = FLASH_CS_PIN; _delay_us(10);

#define FLASH_MAGIC 0xd309

#define PART_TYPE_UNUSED 0x0
#define PART_TYPE_DEFIMAGE 0x01
#define PART_TYPE_BACKUPIMAGE 0x02
#define PART_TYPE_LOG 0x03
#define PART_TYPE_RESERVED 0xfe
#define PART_TYPE_INVALID 0xff

/* structure of partition table we write to the chip */
typedef struct {
		uint32_t start; /* start address */
		uint32_t length; /* length of the block */
		uint8_t type; /* partition type */
} part_entry_t;

typedef struct {
		uint16_t magic; /* magic number to detect partition table and prepared flash */
		part_entry_t part[4]; /* four partition table entries */
} part_table_t;

/* the actual table */
part_table_t partition_table;

uint8_t at25df_good, devid1, devid2, manu;
uint16_t page_size; /* how big are pages on the device */
uint32_t max_address; /* end address */
uint32_t part_location; /* where we will find our magic, probably in a boot block */

/* prototypes */
uint8_t at25df_part_detect(void);
uint8_t at25df_write_page(uint32_t address, uint16_t len, uint8_t *source);
uint8_t at25df_erase_block_4k(uint32_t address);
uint8_t at25df_init_parttable(void);
uint8_t at25df_unprotect_chip(void);
uint8_t at25df_erase_chip(void);

/* init the chip, and validate it's identified properly */
void at25df_init(void) {

	/* set the init flag to good, and we'll set any failures to 0 */
	at25df_good = 1;

	/* set DDR for our CS pin */
	FLASH_CS_PORT.DIRSET = FLASH_CS_PIN;

	SPI_END; /* ensure SPI pin is high for now, and it gets a clean falling edge */
    _delay_us(50);

    spi_fast_mode0();

	/* send request for device IDs */
	SPI_START;
	spi_txrx(0x9f);
	/* and read result */
	manu = spi_txrx(0);
	devid1 = spi_txrx(0);
	devid2 = spi_txrx(0);
	SPI_END;

	/* report what we get back, we're expecting values as in
    table 11-2, pg27 */

	switch (manu) {
		case 0x1f:
			/* ATMEL */
			serial0_tx_PGM(PSTR("at25df: found Atmel "));
			break;
		default:
			/* unknown */
			serial0_tx_PGM(PSTR("at25df: unknown chip or no response\r\n"));
			at25df_good = 0;
			break;
	}

	if (!at25df_good) {
		return;
	}

	switch (devid1) {
		case 0x44:
			serial0_tx_PGM(PSTR("AT25/26DFxxx 4-Mbit\r\n"));
			page_size = 256;
			max_address = 0x7ffff;
			part_location = 0x78000;
			break;
		default:
			serial0_tx_PGM(PSTR("unknown\r\n"));
			at25df_good = 0;
			break;
	}

	if (!at25df_good) {
			return;
	}

	/* confirm we have prepared the chip before */
	if (!at25df_part_detect()) {
		at25df_good = 0;
	}

	return;
}

/* return the detected flash chip's page size */
uint16_t at25df_page_size(void) {
	if (at25df_good) {
		return page_size;
	}
	return 0;
}

/* read from the flash starting from the given address, a len length, and write to target */
uint8_t at25df_read_array(uint32_t address, uint16_t len, uint8_t *target) {

	if (!at25df_good) {
		serial0_tx_PGM(PSTR("at25df: declining to read from missing chip\r\n"));
		return 0;
	}

	if (address > max_address) {
		serial0_tx_PGM(PSTR("at25df: attempt to start read beyond end of chip\r\n"));
		return 0;
	}

    spi_fast_mode0();

	/* select the chip */
	SPI_START;

	/* request the read */
	spi_txrx(AT25_READ_ARRAY);
	spi_txrx((address & 0xff0000) >> 16); /* Address MSB */
	spi_txrx((address & 0xff00) >> 8);
	spi_txrx(address & 0xff); /* Address LSB */
	spi_txrx(0); /* don't care byte according to 7.1 pg 8 */

	/* bulk read the data */
	while (len) {
		(*target) = spi_txrx(0); /* a byte of data from the flash */
		target++;
		len--;
		_delay_us(10);
	}

	/* done reading, deselect chip */
	SPI_END;

#ifdef DEBUG_AT25DF
	serial0_tx_PGM(PSTR("at25df: read succeeded\r\n"));
#endif

	return 1;
}

/* write a single page, this can loop if you make len > page length, so we fail the write if that's attempted, but can start anywhere in the page */
uint8_t at25df_write_page(uint32_t address, uint16_t len, uint8_t *source) {
	uint8_t status;

	if (!at25df_good) {
		return 0;
	}

	if (len > page_size) {
		serial0_tx_PGM(PSTR("at25df: write length larger than page, rejected\r\n"));
		return 0;
	}

	if (address % page_size != 0) {
		serial0_tx_PGM(PSTR("at25df: warning: not writing on page boundry, write may loop\r\n"));
	}

    spi_fast_mode0();

	SPI_START;
	spi_txrx(AT25_WRITE_ENABLE);
	SPI_END;

	_delay_us(10);

	SPI_START;
	spi_txrx(AT25_READ_STATUS);
	status = spi_txrx(0);
	SPI_END;

	if (!(status & AT25_STATUS_WEL)) {
		serial0_tx_PGM(PSTR("at25df: failed to enter write mode?\r\n"));
		return 0;
	}

	SPI_START;
	spi_txrx(AT25_BYTEPAGE_PROGRAM);
	spi_txrx((address & 0xff0000) >> 16); /* Address MSB */
	spi_txrx((address & 0xff00) >> 8);
	spi_txrx(address & 0xff); /* Address LSB */
	while (len) {
		spi_txrx((*source));
		source++;
		len--;
	}
	SPI_END;

	_delay_us(10);

	SPI_START;
	spi_txrx(AT25_READ_STATUS);
	while (1) {
		status = spi_txrx(0);
		if (!(status & AT25_STATUS_BSY)) {
			break;
		}
		_delay_us(10);
	}
	status = spi_txrx(0);
	SPI_END;

	/* check write status */
	if (status & AT25_STATUS_EPE) {
		serial0_tx_PGM(PSTR("at25df: failed to write to address "));
		serial0_tx_hex(((address+1) & 0xff0000) >> 16);
		serial0_tx_hex(((address+1) & 0xff00) >> 8);
		serial0_tx_hex((address+1) & 0xff);
		serial0_tx_cr();
		return 0;
	}

	return 1;
}

/* write (slowly) to an arbitrary point without regard for pages, given start address, len bytes, and a source */
uint8_t at25df_write_array(uint32_t address, uint16_t len, uint8_t *source) {
	uint8_t status;

	if (!at25df_good) {
		serial0_tx_PGM(PSTR("at25df: declining to write to missing chip\r\n"));
		return 0;
	}

	if (address > max_address) {
		serial0_tx_PGM(PSTR("at25df: attempt to start write beyond end of chip\r\n"));
		return 0;
	}

	/* 1. Write Enable
	 * 2. Seq Prog Mode opcode + address + 1 byte of data
	 * 3. Poll Status
	 * 4. Seq Prog Mode opcode + 1 byte of data
	 * 5. Poll Status
	 * 6. ...
	 * 7. Write Disable
	 */

    spi_fast_mode0();

	/* Write enable */
	SPI_START;
	spi_txrx(AT25_WRITE_ENABLE);
	SPI_END;

	/* Seq Prog Mode + address + 1 data byte */
	SPI_START;
	spi_txrx(AT25_SEQUENCE_PROGRAM_A);
	spi_txrx((address & 0xff0000) >> 16); /* Address MSB */
	spi_txrx((address & 0xff00) >> 8);
	spi_txrx(address & 0xff); /* Address LSB */
	spi_txrx(*source); /* first source byte */
	SPI_END;
	PORTC.OUTTGL = PIN2_bm;
	len--;
	source++;

	/* wait for confirmation the byte wrote */
	SPI_START;
	spi_txrx(AT25_READ_STATUS);
	while (1) {
		status = spi_txrx(0);
		if (!(status & AT25_STATUS_BSY)) {
			break;
		}
	}
	SPI_END;

	/* check write status */
	if (status & AT25_STATUS_EPE) {
		serial0_tx_PGM(PSTR("at25df: failed to write to address "));
		serial0_tx_hex(((address+1) & 0xff0000) >> 16);
		serial0_tx_hex(((address+1) & 0xff00) >> 8);
		serial0_tx_hex((address+1) & 0xff);
		serial0_tx_cr();
		return 0;
	}

	/* remaining bytes *if* we have them */
	if (len) {
		while (len) {
			SPI_START;
			spi_txrx(AT25_SEQUENCE_PROGRAM_A);
			spi_txrx(*source);
			SPI_END;
			source++;
			len--;

			/* wait for confirmation the byte wrote */
			SPI_START;
			spi_txrx(AT25_READ_STATUS);
			while (1) {
				status = spi_txrx(0);
				if (!(status & AT25_STATUS_BSY)) {
					break;
				}
			}
			status = spi_txrx(0);
			SPI_END;

			/* check write status */
			if (status & AT25_STATUS_EPE) {
				serial0_tx_PGM(PSTR("at25df: failed to write to address "));
				serial0_tx_hex(((address+1) & 0xff0000) >> 16);
				serial0_tx_hex(((address+1) & 0xff00) >> 8);
				serial0_tx_hex((address+1) & 0xff);
				serial0_tx_cr();
				return 0;
			}
		}
	}

	/* disable write, to complete the session */
	SPI_START;
	spi_txrx(AT25_WRITE_DISABLE);
	SPI_END;

#ifdef DEBUG_AT25DF
	serial0_tx_PGM(PSTR("at25df: write completed\r\n"));
#endif

	return 1;
}

uint8_t at25df_erase_chip(void) {
	uint8_t status;

    spi_fast_mode0();

	/* this should only be called on a chip which does not contain the correct magic number */
	SPI_START;
	spi_txrx(AT25_WRITE_ENABLE);
	SPI_END;
	_delay_us(10);
	SPI_START;
	spi_txrx(AT25_CHIP_ERASE_A);
	SPI_END;
	_delay_us(10);
	/* wait for confirmation the erase is complete */
	SPI_START;
	spi_txrx(AT25_READ_STATUS);
	while (1) {
		status = spi_txrx(0);
		if (!(status & AT25_STATUS_BSY)) {
			break;
		}
		_delay_us(10);
	}
	status = spi_txrx(0);
	SPI_END;

	if (status & AT25_STATUS_EPE) {
		serial0_tx_PGM(PSTR("at25df: unable to erase chip, write protected?\r\n"));
		return 0;
	}

	return 1;
}

uint8_t at25df_unprotect_chip(void) {
	uint8_t status;

    spi_fast_mode0();

	/* first enable writes */
	SPI_START;
	spi_txrx(AT25_WRITE_ENABLE);
	SPI_END;

	_delay_us(10);

	/* now issue global unprotect. this is overloaded on status register */
	SPI_START;
	spi_txrx(AT25_WRITE_STATUS);
	spi_txrx(0x0); /* global unprotect */
	SPI_END;

	/* check lock bits are now clear */
	SPI_START;
	spi_txrx(AT25_READ_STATUS);
	status = spi_txrx(0);
	SPI_END;

	if (status & (AT25_STATUS_SWP0 | AT25_STATUS_SWP1)) {
		serial0_tx_PGM(PSTR("at25df: unable to unprotect chip, write protected?\r\n"));
		return 0;
	}

	return 1;
}

uint8_t at25df_part_detect(void) {

	if (!at25df_good) {
		return 0;
	}

	/* read from where the paritition table should be */
	at25df_read_array(part_location,sizeof(part_table_t),(uint8_t *)&partition_table);
	if (partition_table.magic == FLASH_MAGIC) {
		serial0_tx_PGM(PSTR("at25df: found valid partition table\r\n"));
		return 0;
	}

	serial0_tx_PGM(PSTR("at25df: no partition table found, erasing chip\r\n"));

	/* unprotect the chip so full chip erase will succeed */
	if (!at25df_unprotect_chip()) {
		return 0;
	}

	/* now erase it */
	if (!at25df_erase_chip()) {
		return 0;
	}

	/* write new table */
	if (!at25df_init_parttable()) {
		return 0;
	}

	serial0_tx_PGM(PSTR("at25df: partition table written\r\n"));

	return 1;

}

uint8_t at25df_erase_block_4k(uint32_t address) {
	uint8_t status;

    spi_fast_mode0();

	/* first enable writing to the chip */
	SPI_START;
	spi_txrx(AT25_WRITE_ENABLE);
	SPI_END;

	_delay_us(10);

	/* then issue the block erase command */
	SPI_START;
	spi_txrx(AT25_BLOCK_ERASE_4K);
	spi_txrx((address & 0xff0000) >> 16); /* Address MSB */
	spi_txrx((address & 0xff00) >> 8);
	spi_txrx(address & 0xff); /* Address LSB */
	SPI_END;

	_delay_us(10);

	/* check erase status */
	SPI_START;
	spi_txrx(AT25_READ_STATUS);
	while (1) {
		status = spi_txrx(0);
		if (!(status & AT25_STATUS_BSY)) {
			break;
		}
		_delay_us(10);
	}
	status = spi_txrx(0);
	SPI_END;

	if (status & AT25_STATUS_EPE) {
		serial0_tx_PGM(PSTR("at25df: failed to erase block\r\n"));
		return 0;
	}

	return 1;
}

/* initalise the partition table in local memory, and write it out */
uint8_t at25df_init_parttable(void) {
	/* clear it first */
	memset(&partition_table,0,sizeof(part_table_t));

	/* make sure it has magic number for detection */
	partition_table.magic = FLASH_MAGIC;

	/* partition according to detected chip */
	switch (manu) {
		case 0x1f: /* Atmel */
			switch (devid1) {
				case 0x44: /* 4-Mbit */
					partition_table.part[0].type = PART_TYPE_DEFIMAGE;
					partition_table.part[0].start = 0x0;
					partition_table.part[0].length = 0x30000; /* 192k */
					partition_table.part[1].type = PART_TYPE_BACKUPIMAGE;
					partition_table.part[1].start = 0x30000;
					partition_table.part[1].length = 0x30000; /* 192k */
					partition_table.part[2].type = PART_TYPE_LOG;
					partition_table.part[2].start = 0x60000;
					partition_table.part[2].length = 0x18000; /* 96k */
					partition_table.part[3].type = PART_TYPE_RESERVED;
					partition_table.part[3].start = 0x78000;
					partition_table.part[3].length = 0x8000; /* 32kb boot blocks */
					break;
			}
			break;
	}

	if (!at25df_erase_block_4k(part_location)) {
		return 0;
	}

	/* write the partition table to the block */
	if (!at25df_write_page(part_location, sizeof(part_table_t),(uint8_t *)&partition_table)) {
		serial0_tx_PGM(PSTR("at25df: failed to write new partition table\r\n"));
		return 0;
	}

	/* check it's written */
	memset(&partition_table,0,sizeof(part_table_t));
	at25df_read_array(part_location,sizeof(part_table_t),(uint8_t *)&partition_table);
	if (partition_table.magic != FLASH_MAGIC) {
		serial0_tx_PGM(PSTR("at25df: failed to write new partition table\r\n"));
		return 0;
	}

	return 1;
}
